# train_and_eval_yolo.py
from ultralytics import YOLO
import os
# 设置PyTorch内存分配策略
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'


def train_yolo(data='mydata.yaml', model='yolov8n.pt', epochs=100, imgsz=640, device=0):
    # 加载模型
    model = YOLO(model)
    # 训练
    # model.train(data=data, epochs=epochs, imgsz=imgsz, device=device,batch=8)

    # 优化后的训练参数
    model.train(
        data="mydata.yaml",
        epochs=150,             # 增加训练轮次
        imgsz=640,
        device="0",
        batch=8,
        # patience=30,            # 减少早停耐心值
        # cache="disk",            # 使用RAM缓存加速训练
        # optimizer="Adam",       # 切换到Adam优化器
        # cos_lr=True,            # 启用余弦学习率调度
        # close_mosaic=20,        # 最后20轮关闭Mosaic增强
        # dropout=0.1,            # 添加dropout防止过拟合
        # multi_scale=True,       # 启用多尺度训练
        # lr0=0.001,              # 降低初始学习率
        # lrf=0.1,                # 增大学习率衰减率
        # augment=True,           # 启用数据增强
        hsv_h=0.015,            # 色调扰动
        hsv_s=0.7,              # 饱和度扰动
        hsv_v=0.4,              # 亮度扰动
        fliplr=0.5,             # 水平翻转
        # mosaic=0.8,             # Mosaic增强 (降低至0.8)
        # mixup=0.1,              # 启用MixUp增强
        # conf=0.25,              # 设置置信度阈值
        # iou=0.6,                # 降低NMS的IOU阈值
        # max_det=150,            # 减少最大检测数
        # agnostic_nms=True,      # 类别无关NMS
        # save_txt=True,          # 保存预测结果为文本
        # save_conf=True,         # 保存置信度分数
        # save_crop=True          # 保存裁剪的检测框
    )




def eval_yolo(model_path='runs/detect/train/weights/best.pt', data='mydata.yaml', device=0):
    model = YOLO(model_path)
    metrics = model.val(data=data, device=device)
    print(metrics)

if __name__ == "__main__":
    # 训练
    train_yolo()
    # 测试
    eval_yolo()